from pypdf import PdfReader, PdfWriter, constants
from reportlab.pdfgen import canvas
from reportlab.lib.pagesizes import A4
from reportlab.pdfbase.pdfmetrics import stringWidth
from PIL import Image
import os
import io
import argparse
import uuid
import random
import string

def get_parse():
    parse = argparse.ArgumentParser(description=\
'''
Description:
    pdf tools, add watermark and page number.
Usage:
    pdf_add_watermark_pagination.py -img watermark.png -i input.pdf -o output.pdf               # add watermark, not add page number
    pdf_add_watermark_pagination.py -img watermark.png -i input.pdf -o output.pdf -p 0 0 1      # also add page number, no cover, no informal, first page of formal is p1.
    pdf_add_watermark_pagination.py -img watermark.png -i input.pdf -o output.pdf -p 1 1 4      # also add page number, has cover, has informal, first page of formal is p4.
'''
, formatter_class=argparse.RawTextHelpFormatter)
    parse.add_argument("-img", "--image", type=str, default="maple_confidential.png", metavar="xxx.png", help="input watermark image, png format best")
    parse.add_argument("-p" , "--page_number_info", type=int, default=[], nargs=3, metavar="0", help="add page number, [has_cover, has_informal, first_page_of_formal]")
    parse.add_argument("-i" , "--input", type=str, default="input.pdf", metavar="xxx.pdf", help="input pdf file path")
    parse.add_argument("-o" , "--output", type=str, default="output.pdf", metavar="xxx.pdf", help="output pdf file path")

    return parse.parse_args()

def int_to_roman(num):
    val = [
        1000, 900, 500, 400,
        100, 90, 50, 40,
        10, 9, 5, 4,
        1
        ]
    syb = [
        "M", "CM", "D", "CD",
        "C", "XC", "L", "XL",
        "X", "IX", "V", "IV",
        "I"
        ]
    roman_num = ''
    i = 0
    while num > 0:
        for _ in range(num // val[i]):
            roman_num += syb[i]
            num -= val[i]
        i += 1
    return roman_num


class pdf_release:
    def __init__(self, input_pdf_path, output_pdf_path, watermark_png_path, zoom:float, transparency:float, page_number_info=None):
        self.input_pdf_path = input_pdf_path
        self.output_pdf_path = output_pdf_path
        self.watermark_png_path = watermark_png_path
        self.zoom = zoom
        self.transparency = transparency
        self.temp_watermark_pdf_path = "watermark_temp.pdf"
        self.page_number_info = page_number_info
        if self.page_number_info:
            self.has_cover = self.page_number_info[0]
            self.has_informal = self.page_number_info[1]
            self.first_page_of_formal = self.page_number_info[2]
            self.page_number_info_check()
        else:
            self.has_cover = None
            self.has_informal = None
            self.first_page_of_formal = None

    def page_number_info_check(self):
        if not self.has_cover and not self.has_informal:    # [0, 0, 1]
            if self.first_page_of_formal != 1:
                print("\033[33mWarning:\033[0m first_page_of_formal must be 1 when no cover and no informal.")
                self.first_page_of_formal = 1
        elif not self.has_cover and self.has_informal:    # [0, 1, >=2]
            if self.first_page_of_formal < 2:
                print("\033[33mWarning:\033[0m first_page_of_formal must be >=2 when no cover and has informal.")
                self.first_page_of_formal = 2
        elif self.has_cover and not self.has_informal:    # [1, 0, 2]
            if self.first_page_of_formal != 2:
                print("\033[33mWarning:\033[0m first_page_of_formal must be 2 when has cover and no informal.")
                self.first_page_of_formal = 2
        elif self.has_cover and self.has_informal:    # [1, 1, >=3]
            if self.first_page_of_formal < 3:
                print("\033[33mWarning:\033[0m first_page_of_formal must be >=3 when has cover and has informal.")
                self.first_page_of_formal = 3
            
    
    def create_watermark_pdf(self):
        # 创建一个临时的水印PDF文件
        c = canvas.Canvas(self.temp_watermark_pdf_path, pagesize=A4)
        c.setFillAlpha(self.transparency)  # 设置水印透明度

        # 获取水印图片的尺寸
        watermark_img = Image.open(self.watermark_png_path)
        watermark_width, watermark_height = watermark_img.size
        # 缩放到50%
        watermark_width *= zoom
        watermark_height *= zoom

        # 获取PDF页面的尺寸
        self.pdf_reader = PdfReader(self.input_pdf_path)
        first_page = self.pdf_reader.pages[0]
        page_width = first_page.mediabox[2]
        page_height = first_page.mediabox[3]

        # 计算水印居中的位置
        x = (float(page_width) - watermark_width) / 2
        y = (float(page_height) - watermark_height) / 2

        # 旋转画布45度
        c.saveState()  # 保存当前状态，以便之后恢复
        c.translate(float(page_width) / 2, float(page_height) / 2)  # 将旋转中心移动到页面中心
        c.rotate(45)  # 旋转45度
        # 添加水印图片，位置需要重新计算以反映旋转
        c.drawImage(watermark_png_path, -watermark_width / 2, -watermark_height / 2, width=watermark_width, height=watermark_height, mask='auto')
        c.restoreState()  # 恢复到之前的状态

        c.save()

    @staticmethod
    def delete_temp_pdf(temp_pdf_path):
        # 删除临时的水印PDF文件
        if os.path.exists(temp_pdf_path):
            os.remove(temp_pdf_path)


    @staticmethod
    def create_page_number_pdf(current_page, total_pages, is_roman=False):
        packet = io.BytesIO()
        c = canvas.Canvas(packet, pagesize=A4)

        if not is_roman:
            page_number_text = f"{current_page}/{total_pages}"
        else:
            # 将当前页码和总页数转换为罗马数字
            current_page_roman = int_to_roman(current_page)
            total_pages_roman = int_to_roman(total_pages)
            page_number_text = f"{current_page_roman}/{total_pages_roman}"
        
        # 设置字体和大小
        font_name = "Times-Roman" # 字体,有"Courier", "Helvetica", "Times-Roman", "Symbol", "ZapfDingbats"
        font_size = 10.5 # 10.5 五号, 12 小四
        c.setFont(font_name, font_size)
        
        # 计算页码字符串宽度
        text_width = stringWidth(page_number_text, font_name, font_size)
        
        # A4纸宽度的一半减去字符串宽度的一半，以实现居中
        x_position = (A4[0] - text_width) / 2
        y_position = 20  # 距离底部20点的高度
        
        c.drawString(x_position, y_position, page_number_text)
        c.save()
        packet.seek(0)
        return packet
        
        # 保存到文件
        # with open(output_path, 'wb') as f:
        #     f.write(packet.getvalue())

    @staticmethod
    def generate_random_string(bits=128):
        characters = string.ascii_letters + string.digits  # 包含小写字母、大写字母和数字
        random_string = ''.join(random.choices(characters, k=bits))  # 随机选择128个字符
        return random_string
    

    def create_final_pdf(self):
        # 合并水印PDF和原始PDF
        self.watermark_page = PdfReader(self.temp_watermark_pdf_path).pages[0]
        self.pdf_writer = PdfWriter()

        self.pdf_writer.clone_document_from_reader(self.pdf_reader)
        # self.pdf_writer.clone_reader_document_root(self.pdf_reader)

        total_pages = len(self.pdf_writer.pages)  # 计算总页数

        for i, page in enumerate(self.pdf_writer.pages):
            page.merge_page(self.watermark_page)  # 合并水印到当前页
            if self.page_number_info:
                if i==0 and self.has_cover:
                    continue
                elif i<self.first_page_of_formal-1 and self.has_informal:
                    if self.has_cover:
                        page_number_page = PdfReader(self.create_page_number_pdf(i, self.first_page_of_formal-2, is_roman=True)).pages[0]
                    else:
                        page_number_page = PdfReader(self.create_page_number_pdf(i+1, self.first_page_of_formal-2+1, is_roman=True)).pages[0]
                else:
                    page_number_page = PdfReader(self.create_page_number_pdf(i+1-self.first_page_of_formal+1, total_pages-self.first_page_of_formal+1, is_roman=False)).pages[0]
                page.merge_page(page_number_page) # 合并页码到当前页
        
        # 加密，限制编辑
        edit_passwoord = self.generate_random_string(bits=128)
        self.pdf_writer.encrypt(user_password="", owner_password=edit_passwoord, use_128bit=True, permissions_flag=constants.UserAccessPermissions.ADD_OR_MODIFY)
        print("\033[34mInfo:\033[0m The output pdf file is encrypted, the edit password is:", edit_passwoord)

        with open(output_pdf_path, 'wb') as output_pdf_file:
            self.pdf_writer.write(output_pdf_file)


if __name__ == '__main__':

    args = get_parse()
    watermark_png_path = args.image
    input_pdf_path = args.input
    output_pdf_path = args.output
    page_number_info = args.page_number_info
    # page_number_info = [1, 1, 4] # 1:有封面; 1:有非正文内容; 4:正文是从整体的第4页开始

    zoom = 0.5
    transparency = 0.2

    pdf_obj = pdf_release(input_pdf_path, output_pdf_path, watermark_png_path, zoom, transparency, page_number_info)
    pdf_obj.create_watermark_pdf()
    pdf_obj.create_final_pdf()
    pdf_obj.delete_temp_pdf(pdf_obj.temp_watermark_pdf_path)